#!/usr/bin/env python3
"""
快速热更新测试脚本

简单测试YOLO模型热更新功能
"""

import requests
import time
import json
import argparse

def test_hot_swap(server_url: str = "http://localhost:8000"):
    """测试热更新功能"""
    
    print("🔥 YOLO模型热更新测试")
    print("=" * 50)
    
    # 1. 检查服务器状态
    print("1. 检查服务器状态...")
    try:
        response = requests.get(f"{server_url}/health", timeout=5)
        if response.status_code == 200:
            print("✅ 服务器运行正常")
        else:
            print(f"❌ 服务器状态异常: {response.status_code}")
            return
    except Exception as e:
        print(f"❌ 无法连接服务器: {e}")
        return
    
    # 2. 获取可用模型列表
    print("\n2. 获取模型列表...")
    try:
        response = requests.get(f"{server_url}/models")
        if response.status_code == 200:
            models_data = response.json()
            models = models_data.get("models", [])
            print(f"✅ 找到 {len(models)} 个模型:")
            for model in models:
                status = "🔥活跃" if model.get("active") else ("✅已加载" if model.get("loaded") else "⏳未加载")
                print(f"   - {model['name']}: {status}")
        else:
            print(f"❌ 获取模型列表失败: {response.status_code}")
            return
    except Exception as e:
        print(f"❌ 获取模型列表失败: {e}")
        return
    
    # 3. 选择测试模型
    available_models = [m["name"] for m in models if m.get("available", True)]
    if len(available_models) < 2:
        print("❌ 需要至少2个可用模型进行热更新测试")
        return
    
    model1, model2 = available_models[0], available_models[1]
    print(f"\n3. 测试模型: {model1} ↔ {model2}")
    
    # 4. 测试热更新
    test_models = [model1, model2, model1]  # 来回切换
    
    for i, target_model in enumerate(test_models, 1):
        print(f"\n4.{i} 热更新到模型: {target_model}")
        
        # 获取当前状态
        try:
            response = requests.get(f"{server_url}/models/{target_model}/status")
            if response.status_code == 200:
                status = response.json()
                print(f"   当前状态: {'🔥活跃' if status.get('active') else '⏳未活跃'}")
            else:
                print(f"   无法获取状态: {response.status_code}")
        except Exception as e:
            print(f"   状态查询失败: {e}")
        
        # 执行热更新
        try:
            start_time = time.time()
            response = requests.post(f"{server_url}/models/{target_model}/hot-swap", timeout=30)
            swap_time = time.time() - start_time
            
            if response.status_code == 200:
                result = response.json()
                server_time = result.get("swap_time", swap_time)
                print(f"   ✅ 热更新成功!")
                print(f"   📊 服务器耗时: {server_time:.3f}s")
                print(f"   📊 客户端耗时: {swap_time:.3f}s")
                
                # 验证状态
                time.sleep(0.5)  # 等待状态更新
                response = requests.get(f"{server_url}/models/{target_model}/status")
                if response.status_code == 200:
                    status = response.json()
                    if status.get("active"):
                        print(f"   ✅ 状态验证: 模型已激活")
                    else:
                        print(f"   ⚠️ 状态验证: 模型未激活")
                
            else:
                error_text = response.text
                print(f"   ❌ 热更新失败: {response.status_code}")
                print(f"   错误信息: {error_text}")
                
        except Exception as e:
            print(f"   ❌ 热更新异常: {e}")
        
        # 等待一下再进行下一次测试
        if i < len(test_models):
            time.sleep(2)
    
    # 5. 最终状态检查
    print(f"\n5. 最终状态检查...")
    try:
        response = requests.get(f"{server_url}/models")
        if response.status_code == 200:
            models_data = response.json()
            models = models_data.get("models", [])
            print("📊 最终模型状态:")
            for model in models:
                if model["name"] in available_models:
                    status = "🔥活跃" if model.get("active") else ("✅已加载" if model.get("loaded") else "⏳未加载")
                    usage = model.get("usage_count", 0)
                    print(f"   - {model['name']}: {status} (使用次数: {usage})")
    except Exception as e:
        print(f"❌ 最终状态检查失败: {e}")
    
    print("\n🎉 热更新测试完成!")

def main():
    parser = argparse.ArgumentParser(description="YOLO模型热更新快速测试")
    parser.add_argument("--server", default="http://localhost:8000", help="服务器地址")
    
    args = parser.parse_args()
    
    test_hot_swap(args.server)

if __name__ == "__main__":
    main()
